import csv
import os
from django.core.management.base import BaseCommand
from django.contrib.auth.models import User
from yourapp.models import Movie, Rating

class Command(BaseCommand):
    help = 'Import MovieLens dataset'

    def add_arguments(self, parser):
        parser.add_argument('dataset_path', type=str, help='Path to MovieLens dataset directory')

    def handle(self, *args, **options):
        dataset_path = options['dataset_path']
        
        # 导入电影数据
        movies_file = os.path.join(dataset_path, 'movies.csv')
        self.stdout.write(f'Importing movies from {movies_file}...')
        
        with open(movies_file, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            next(reader)  # 跳过标题行
            
            for row in reader:
                movie_id, title, genres = row
                Movie.objects.create(
                    movie_id=int(movie_id),
                    title=title,
                    genres=genres
                )
        
        self.stdout.write(self.style.SUCCESS('Successfully imported movies'))
        
        # 导入评分数据 (可能非常大，所以分批处理)
        ratings_file = os.path.join(dataset_path, 'ratings.csv')
        self.stdout.write(f'Importing ratings from {ratings_file}...')
        
        # 确保有足够的用户
        # 这里为简单起见，我们创建与数据集中用户ID对应的User对象
        # 实际应用中，您可能需要将MovieLens用户ID映射到您系统中的用户
        
        with open(ratings_file, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            next(reader)  # 跳过标题行
            
            batch_size = 10000
            batch = []
            
            for i, row in enumerate(reader):
                user_id, movie_id, rating, timestamp = row
                
                # 确保用户存在
                user, created = User.objects.get_or_create(
                    username=f'user_{user_id}',
                    defaults={'password': 'movielens'}
                )
                
                # 获取电影
                try:
                    movie = Movie.objects.get(movie_id=int(movie_id))
                    
                    # 添加到批次
                    batch.append(Rating(
                        user=user,
                        movie=movie,
                        rating=float(rating),
                        timestamp=int(timestamp)
                    ))
                    
                    # 达到批次大小时批量创建
                    if len(batch) >= batch_size:
                        Rating.objects.bulk_create(batch, ignore_conflicts=True)
                        batch = []
                        self.stdout.write(f'Imported {i+1} ratings...')
                        
                except Movie.DoesNotExist:
                    continue
            
            # 导入剩余的批次
            if batch:
                Rating.objects.bulk_create(batch, ignore_conflicts=True)
                
        self.stdout.write(self.style.SUCCESS('Successfully imported ratings')) 